#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/01/12 

# ARIP data, 分析 Kind 和 Pair 之间的正向/反向概率分布

from pathlib import Path
from collections import Counter
from typing import List

import pandas as pd
import matplotlib.pyplot as plt ; plt.ion()

from data import RDATA_PATH

DATA_FILE = 'InterChain_41D.csv'

def plot(dfs:List[pd.DataFrame], target, titles=[], suptitle=''):
  n_plot = len(dfs)
  plt.clf()
  for i in range(n_plot):
    plt.subplot(n_plot, 1, i+1)
    cnt_pair = sorted([(c, v) for v, c in Counter(dfs[i][target]).items()], reverse=True)
    freqs = [c for c, v in cnt_pair]
    pairs = [v for c, v in cnt_pair]
    plt.xticks(range(len(pairs)), pairs)
    plt.plot(freqs, label=titles[i] if titles else None)
    plt.legend()
  plt.suptitle(suptitle)
  plt.show()


def prob_kind_pair():
  df = pd.read_csv(Path(RDATA_PATH) / DATA_FILE)
  print(f'>> len(df): {len(df)}')
  columns = ['Prop', 'Kind', 'Pair']
  df = df[columns]
  print('=' * 72)

  kinds = sorted(set(df['Kind']))
  print(f'>> kinds: {kinds}')
  residues = sorted({ r for pair in set(df['Pair']) for r in pair.split('-') })
  print(f'>> residues: {residues}')
  pairs = set(df['Pair'])

  print('=' * 72)
  while True:
    try:
      s = input('>> input <kind> or <residue>-<residue>: ').strip()
      
      if s == 'q': break
      elif s in kinds:
        sdf = df[df['Kind'] == s]
        sdf_s = sdf[sdf['Prop'] == 'Stable']
        sdf_f = sdf[sdf['Prop'] == 'Flexible']
        plot([sdf, sdf_s, sdf_f], 'Pair', titles=['all', 'stab', 'flex'], suptitle=s)
      elif s in pairs:
        sdf = df[df['Pair'] == s]
        sdf_s = sdf[sdf['Prop'] == 'Stable']
        sdf_f = sdf[sdf['Prop'] == 'Flexible']
        plot([sdf, sdf_s, sdf_f], 'Kind', titles=['all', 'stab', 'flex'], suptitle=s)
      else: print('<< unrecognized input') ; continue
    except KeyboardInterrupt: break


if __name__ == '__main__':
  prob_kind_pair()
